package com.lm.langchain4j.adapter;


import dev.langchain4j.model.embedding.EmbeddingModel;

import java.lang.reflect.Method;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 反射尝试调用 embedding 模型的常见方法签名，
 * 并将结果标准化为 List<double[]>。
 *
 * 兼容常见返回类型：
 *  - List<Embedding>（Embedding 中有 vector 或 values）
 *  - List<List<Float>> / List<List<Double>>
 *  - List<float[]> / List<double[]>
 */
public class EmbeddingModelAdapter {

    private final EmbeddingModel embeddingModel;
    private final Class<?> modelClass;
    private final List<Method> candidateMethods = new ArrayList<>();

    public EmbeddingModelAdapter(EmbeddingModel  embeddingModel) {
        this.embeddingModel = Objects.requireNonNull(embeddingModel);
        this.modelClass = embeddingModel.getClass();
        initCandidateMethods();
    }

    private void initCandidateMethods() {
        // 常见方法名： embedAll, embedTexts, embed, embedBatch, embeddings
        String[] names = new String[]{"embedAll", "embedTexts", "embed", "embedMany", "embeddings"};
        for (Method m : modelClass.getMethods()) {
            for (String n : names) {
                if (m.getName().equalsIgnoreCase(n)) {
                    candidateMethods.add(m);
                }
            }
        }
        // 也尝试单字符串签名
        for (Method m : modelClass.getMethods()) {
            if (m.getName().equalsIgnoreCase("embed") && m.getParameterCount() == 1 &&
                    m.getParameterTypes()[0] == String.class) {
                candidateMethods.add(m);
            }
        }
    }

    /**
     * 将 texts 向量化，返回 List<double[]>，顺序对应 texts。
     */
    public List<double[]> embedTexts(List<String> texts) {
        if (texts == null || texts.isEmpty()) return Collections.emptyList();

        // 先尝试 batch 方法
        for (Method m : candidateMethods) {
            try {
                Object result = tryInvoke(m, texts);
                List<double[]> vectors = normalizeResult(result);
                if (vectors != null && vectors.size() == texts.size()) return vectors;
            } catch (Throwable ignored) {
            }
        }

        // 退化：对每个文本单独调用 embed(String)
        try {
            Method single = modelClass.getMethod("embed", String.class);
            List<double[]> out = new ArrayList<>();
            for (String t : texts) {
                Object r = single.invoke(embeddingModel, t);
                List<double[]> norm = normalizeResult(r);
                if (norm != null && !norm.isEmpty()) out.add(norm.get(0));
                else throw new IllegalStateException("single embed returned null");
            }
            return out;
        } catch (Throwable ex) {
            throw new RuntimeException("无法调用 embedding 方法，请检查 embeddingModel API", ex);
        }
    }

    private Object tryInvoke(Method m, List<String> texts) throws Exception {
        m.setAccessible(true);
        Class<?>[] pts = m.getParameterTypes();
        if (pts.length == 1) {
            Class<?> p = pts[0];
            if (List.class.isAssignableFrom(p)) {
                return m.invoke(embeddingModel, texts);
            } else if (p == String[].class) {
                return m.invoke(embeddingModel, (Object) texts.toArray(new String[0]));
            } else if (p == String.class) {
                // single string caller (shouldn't reach here in batch)
                return m.invoke(embeddingModel, String.join("\n", texts));
            }
        } else if (pts.length == 2) {
            // some APIs accept (List<String>, Map<String,?>) options
            if (List.class.isAssignableFrom(pts[0])) {
                return m.invoke(embeddingModel, texts, Collections.emptyMap());
            }
        }
        // unsupported signature
        throw new NoSuchMethodException("不支持的方法签名: " + m);
    }

    /**
     * 将多种可能的返回值统一为 List<double[]>。
     */
    @SuppressWarnings("unchecked")
    private List<double[]> normalizeResult(Object result) {
        if (result == null) return null;

        // case: List<double[]>
        if (result instanceof List) {
            List<?> list = (List<?>) result;
            if (list.isEmpty()) return Collections.emptyList();
            Object first = list.get(0);

            if (first instanceof double[]) {
                return (List<double[]>) list;
            }

            if (first instanceof float[]) {
                return list.stream()
                        .map(o -> {
                            float[] fa = (float[]) o;
                            double[] da = new double[fa.length];
                            for (int i = 0; i < fa.length; i++) da[i] = fa[i];
                            return da;
                        }).collect(Collectors.toList());
            }

            // List<List<Double>> or List<List<Float>>
            if (first instanceof List) {
                return list.stream()
                        .map(o -> {
                            List<?> nums = (List<?>) o;
                            double[] da = new double[nums.size()];
                            for (int i = 0; i < nums.size(); i++) {
                                Object v = nums.get(i);
                                if (v instanceof Number) da[i] = ((Number) v).doubleValue();
                                else da[i] = Double.parseDouble(v.toString());
                            }
                            return da;
                        }).collect(Collectors.toList());
            }

            // List<Embedding> where Embedding has vector/getValues/getVector etc.
            try {
                // try to extract numeric array via reflection
                return list.stream()
                        .map(o -> {
                            try {
                                Class<?> c = o.getClass();
                                // common getter names
                                String[] getters = new String[]{"getVector", "vector", "getValues", "getEmbedding", "getData"};
                                for (String g : getters) {
                                    try {
                                        Method gm = c.getMethod(g);
                                        Object val = gm.invoke(o);
                                        if (val instanceof List) {
                                            List<?> nums = (List<?>) val;
                                            double[] da = new double[nums.size()];
                                            for (int i = 0; i < nums.size(); i++)
                                                da[i] = ((Number) nums.get(i)).doubleValue();
                                            return da;
                                        } else if (val instanceof double[]) {
                                            return (double[]) val;
                                        } else if (val instanceof float[]) {
                                            float[] fa = (float[]) val;
                                            double[] da = new double[fa.length];
                                            for (int i = 0; i < fa.length; i++) da[i] = fa[i];
                                            return da;
                                        }
                                    } catch (NoSuchMethodException ignored) {
                                    }
                                }
                            } catch (Throwable t) {
                            }
                            throw new RuntimeException("无法从 Embedding 对象中提取向量: " + o.getClass());
                        }).collect(Collectors.toList());
            } catch (Throwable t) {
                throw new RuntimeException("normalizeResult 失败", t);
            }
        }

        // unsupported
        return null;
    }
}
